{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FFI DLPack 接口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------\n",
      "Benchmark f(x, y, z) overhead\n",
      "-----------------------------\n",
      "numpy.add                                7.214546203613281e-07 sec/call\n",
      "torch.add[cpu]                           4.431867599487305e-06 sec/call\n",
      "torch.add[cuda]                          1.4083337783813477e-05 sec/call\n",
      "tvm.ffi.nop                              4.0678262710571285e-06 sec/call\n",
      "tvm.ffi.nop+from_dlpack(torch)           3.829042911529541e-05 sec/call\n",
      "tvm.ffi.nop+from_dlpack(numpy)           1.171247959136963e-05 sec/call\n",
      "tvm.ffi.nop+from_dlpack(tvm)             1.4158987998962402e-05 sec/call\n",
      "tvm.ffi.nop+from_dlpack(torch.utils)     2.3821353912353516e-05 sec/call\n",
      "tvm.ffi.nop.autodlpack(torch[cpu])       2.6399731636047363e-05 sec/call\n",
      "tvm.ffi.nop.autodlpack(torch[cuda])      2.7433156967163087e-05 sec/call\n",
      "tvm.ffi.nop.autodlpack(numpy)            1.2148022651672363e-05 sec/call\n",
      "-------------------------------\n",
      "Benchmark x.__dlpack__ overhead\n",
      "-------------------------------\n",
      "torch.utils.dlpack.to_dlpack             1.0176658630371093e-06 sec/call\n",
      "torch.__dlpack__                         3.405594825744629e-06 sec/call\n",
      "numpy.__dlpack__                         1.462697982788086e-07 sec/call\n",
      "tvm.__dlpack__                           5.183458328247071e-07 sec/call\n",
      "---------------------------------------------------\n",
      "Benchmark x.__dlpack__(max_version=(1,1)) overhead\n",
      "---------------------------------------------------\n",
      "torch.__dlpack__(max_version=(1,1))      Tensor.__dlpack__() got an unexpected keyword argument 'max_version'\n",
      "numpy.__dlpack__(max_version=(1,1))      __dlpack__() got an unexpected keyword argument 'max_version'\n",
      "tvm.__dlpack__(max_version=(1,1))        6.298303604125977e-07 sec/call\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from tvm import ffi as tvm_ffi\n",
    "import time\n",
    "\n",
    "\n",
    "def print_speed(name, speed):\n",
    "    print(f\"{name:<40} {speed} sec/call\")\n",
    "\n",
    "\n",
    "def print_error(name, error):\n",
    "    print(f\"{name:<40} {error}\")\n",
    "\n",
    "\n",
    "def baseline_torch_add(repeat):\n",
    "    \"\"\"Run torch.add with one element\"\"\"\n",
    "\n",
    "    def run_bench(device):\n",
    "        x = torch.arange(1, device=device)\n",
    "        y = torch.arange(1, device=device)\n",
    "        z = torch.arange(1, device=device)\n",
    "\n",
    "        torch.add(x, y, out=z)\n",
    "        if device == \"cuda\":\n",
    "            torch.cuda.synchronize()\n",
    "        start = time.time()\n",
    "        for i in range(repeat):\n",
    "            torch.add(x, y, out=z)\n",
    "        # note we deliberately do not use torch.cuda.synchronize()\n",
    "        # because we want to see the overhead of the FFI call.\n",
    "        end = time.time()\n",
    "        print_speed(f\"torch.add[{device}]\", (end - start) / repeat)\n",
    "\n",
    "    # rough take away: add on cuda roughly takes 3e-6 sec/call\n",
    "    run_bench(\"cpu\")\n",
    "    run_bench(\"cuda\")\n",
    "\n",
    "\n",
    "def baseline_numpy_add(repeat):\n",
    "    \"\"\"Run numpy.add with one element\"\"\"\n",
    "    x = np.arange(1)\n",
    "    y = np.arange(1)\n",
    "    z = np.arange(1)\n",
    "\n",
    "    np.add(x, y, out=z)\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        np.add(x, y, out=z)\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(\"numpy.add\", speed)\n",
    "\n",
    "\n",
    "def baseline_cupy_add(repeat):\n",
    "    \"\"\"Run cupy.add with one element\"\"\"\n",
    "    try:\n",
    "        import cupy\n",
    "    except ImportError:\n",
    "        # skip if cupy is not installed\n",
    "        return\n",
    "    x = cupy.arange(1)\n",
    "    y = cupy.arange(1)\n",
    "    z = cupy.arange(1)\n",
    "\n",
    "    cupy.add(x, y, out=z)\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        cupy.add(x, y, out=z)\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(\"cupy.add\", speed)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop(repeat):\n",
    "    \"\"\"Overhead of tvm FFI python call via calling a NOP.\n",
    "\n",
    "    testing.nop is defined in c++ and do nothing.\n",
    "    \"\"\"\n",
    "    nop = tvm_ffi.get_global_func(\"testing.nop\")\n",
    "    x = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    y = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    z = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    nop(x, y, z)\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        y = tvm_ffi.from_dlpack(x)\n",
    "    end = time.time()\n",
    "    print_speed(\"tvm.ffi.nop\", (end - start) / repeat)\n",
    "\n",
    "\n",
    "def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):\n",
    "    \"\"\"run dlpack conversion + tvm.ffi.nop\n",
    "\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    \"\"\"\n",
    "    nop = tvm_ffi.get_global_func(\"testing.nop\")\n",
    "    tx = tvm_ffi.from_dlpack(x)\n",
    "    ty = tvm_ffi.from_dlpack(y)\n",
    "    tz = tvm_ffi.from_dlpack(z)\n",
    "    nop(tx, ty, tz)\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        tx = tvm_ffi.from_dlpack(x)\n",
    "        ty = tvm_ffi.from_dlpack(y)\n",
    "        tz = tvm_ffi.from_dlpack(z)\n",
    "        nop(tx, ty, tz)\n",
    "    end = time.time()\n",
    "    print_speed(name, (end - start) / repeat)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop_from_torch_dlpack(repeat):\n",
    "    \"\"\"run dlpack conversion + tvm.ffi.nop\n",
    "\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    \"\"\"\n",
    "    x = torch.arange(1)\n",
    "    y = torch.arange(1)\n",
    "    z = torch.arange(1)\n",
    "    bench_ffi_nop_from_dlpack(\"tvm.ffi.nop+from_dlpack(torch)\", x, y, z, repeat)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop_from_numpy_dlpack(repeat):\n",
    "    \"\"\"run dlpack conversion + tvm.ffi.nop\n",
    "\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    \"\"\"\n",
    "    x = np.arange(1)\n",
    "    y = np.arange(1)\n",
    "    z = np.arange(1)\n",
    "    bench_ffi_nop_from_dlpack(\"tvm.ffi.nop+from_dlpack(numpy)\", x, y, z, repeat)\n",
    "\n",
    "\n",
    "def tvm_ffi_self_dlpack_nop(repeat):\n",
    "    \"\"\"run dlpack conversion + tvm.ffi.nop\n",
    "\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    \"\"\"\n",
    "    x = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    y = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    z = tvm_ffi.from_dlpack(torch.arange(1))\n",
    "    bench_ffi_nop_from_dlpack(\"tvm.ffi.nop+from_dlpack(tvm)\", x, y, z, repeat)\n",
    "\n",
    "\n",
    "def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):\n",
    "    \"\"\"run dlpack conversion + tvm.ffi.nop\n",
    "\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    \"\"\"\n",
    "    nop = tvm_ffi.get_global_func(\"testing.nop\")\n",
    "    tx = tvm_ffi.from_dlpack(x)\n",
    "    ty = tvm_ffi.from_dlpack(y)\n",
    "    tz = tvm_ffi.from_dlpack(z)\n",
    "    nop(tx, ty, tz)\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        tx = tvm_ffi.from_dlpack(x)\n",
    "        ty = tvm_ffi.from_dlpack(y)\n",
    "        tz = tvm_ffi.from_dlpack(z)\n",
    "        nop(tx, ty, tz)\n",
    "    end = time.time()\n",
    "    print_speed(name, (end - start) / repeat)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop_from_torch_utils_to_dlpack(repeat):\n",
    "    \"\"\"\n",
    "    Measures overhead of running dlpack for each args then invoke\n",
    "    but uses the legacy torch.utils.dlpack.to_dlpack API\n",
    "\n",
    "    This helps to measure possible implementation overhead of torch.\n",
    "    \"\"\"\n",
    "    nop = tvm_ffi.get_global_func(\"testing.nop\")\n",
    "    x = torch.arange(1)\n",
    "    y = torch.arange(1)\n",
    "    z = torch.arange(1)\n",
    "\n",
    "    tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x))\n",
    "    ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y))\n",
    "    tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z))\n",
    "    nop(tx, ty, tz)\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x))\n",
    "        ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y))\n",
    "        tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z))\n",
    "        nop(tx, ty, tz)\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(\"tvm.ffi.nop+from_dlpack(torch.utils)\", speed)\n",
    "\n",
    "\n",
    "def bench_tvm_ffi_nop_autodlpack(name, x, y, z, repeat):\n",
    "    \"\"\"\n",
    "    Measures overhead of running dlpack via auto convert by directly\n",
    "    take torch.Tensor as inputs.\n",
    "    \"\"\"\n",
    "    nop = tvm_ffi.get_global_func(\"testing.nop\")\n",
    "    nop(x, y, z)\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        nop(x, y, z)\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(name, speed)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop_autodlpack_from_torch(repeat, device=\"cpu\"):\n",
    "    \"\"\"\n",
    "    Measures overhead of running dlpack via auto convert by directly\n",
    "    take torch.Tensor as inputs.\n",
    "    \"\"\"\n",
    "    # use larger to ensure alignment req is met\n",
    "    x = torch.arange(1, device=device)\n",
    "    y = torch.arange(1, device=device)\n",
    "    z = torch.arange(1, device=device)\n",
    "    bench_tvm_ffi_nop_autodlpack(f\"tvm.ffi.nop.autodlpack(torch[{device}])\", x, y, z, repeat)\n",
    "\n",
    "\n",
    "def tvm_ffi_nop_autodlpack_from_numpy(repeat):\n",
    "    \"\"\"\n",
    "    Measures overhead of running dlpack via auto convert by directly\n",
    "    take numpy.ndarray as inputs.\n",
    "    \"\"\"\n",
    "    # use larger to ensure alignment req is met\n",
    "    x = np.arange(256)\n",
    "    y = np.arange(256)\n",
    "    z = np.arange(256)\n",
    "    bench_tvm_ffi_nop_autodlpack(\"tvm.ffi.nop.autodlpack(numpy)\", x, y, z, repeat)\n",
    "\n",
    "\n",
    "def bench_to_dlpack(x, name, repeat):\n",
    "    x.__dlpack__()\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        x.__dlpack__()\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(name, speed)\n",
    "\n",
    "\n",
    "def bench_to_dlpack_versioned(x, name, repeat, max_version=(1, 1)):\n",
    "    \"\"\"\n",
    "    Measures overhead of running dlpack with latest 1.1.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        x.__dlpack__(max_version=max_version)\n",
    "        start = time.time()\n",
    "        for i in range(repeat):\n",
    "            x.__dlpack__(max_version=max_version)\n",
    "        end = time.time()\n",
    "        speed = (end - start) / repeat\n",
    "        print_speed(name, speed)\n",
    "    except Exception as e:\n",
    "        print_error(name, e)\n",
    "\n",
    "\n",
    "def bench_torch_utils_to_dlpack(repeat):\n",
    "    \"\"\"\n",
    "    Measures overhead of running torch.utils.dlpack.to_dlpack\n",
    "    \"\"\"\n",
    "    x = torch.arange(1)\n",
    "    torch.utils.dlpack.to_dlpack(x)\n",
    "    start = time.time()\n",
    "    for i in range(repeat):\n",
    "        torch.utils.dlpack.to_dlpack(x)\n",
    "    end = time.time()\n",
    "    speed = (end - start) / repeat\n",
    "    print_speed(\"torch.utils.dlpack.to_dlpack\", speed)\n",
    "\n",
    "\n",
    "def main():\n",
    "    repeat = 10000\n",
    "    print(\"-----------------------------\")\n",
    "    print(\"Benchmark f(x, y, z) overhead\")\n",
    "    print(\"-----------------------------\")\n",
    "    baseline_numpy_add(repeat)\n",
    "    baseline_torch_add(repeat)\n",
    "    baseline_cupy_add(repeat)\n",
    "    tvm_ffi_nop(repeat)\n",
    "    tvm_ffi_nop_from_torch_dlpack(repeat)\n",
    "    tvm_ffi_nop_from_numpy_dlpack(repeat)\n",
    "    tvm_ffi_self_dlpack_nop(repeat)\n",
    "    tvm_ffi_nop_from_torch_utils_to_dlpack(repeat)\n",
    "    tvm_ffi_nop_autodlpack_from_torch(repeat, \"cpu\")\n",
    "    tvm_ffi_nop_autodlpack_from_torch(repeat, \"cuda\")\n",
    "    tvm_ffi_nop_autodlpack_from_numpy(repeat)\n",
    "    print(\"-------------------------------\")\n",
    "    print(\"Benchmark x.__dlpack__ overhead\")\n",
    "    print(\"-------------------------------\")\n",
    "    bench_torch_utils_to_dlpack(repeat)\n",
    "    bench_to_dlpack(torch.arange(1), \"torch.__dlpack__\", repeat)\n",
    "    bench_to_dlpack(np.arange(1), \"numpy.__dlpack__\", repeat)\n",
    "    bench_to_dlpack(tvm_ffi.from_dlpack(torch.arange(1)), \"tvm.__dlpack__\", repeat)\n",
    "    print(\"---------------------------------------------------\")\n",
    "    print(\"Benchmark x.__dlpack__(max_version=(1,1)) overhead\")\n",
    "    print(\"---------------------------------------------------\")\n",
    "    bench_to_dlpack_versioned(torch.arange(1), \"torch.__dlpack__(max_version=(1,1))\", repeat)\n",
    "    bench_to_dlpack_versioned(np.arange(1), \"numpy.__dlpack__(max_version=(1,1))\", repeat)\n",
    "    bench_to_dlpack_versioned(\n",
    "        tvm_ffi.from_dlpack(torch.arange(1)), \"tvm.__dlpack__(max_version=(1,1))\", repeat\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
